#include "AST.h"
#include "TypeSystem.h"
#include "Scope.h"

//We have computed First+(<exp>) :{'(',<const>,<ident>} by hand
#define isExpression(tok) ((tok) == '(' || (tok) == tok_number || (tok) == tok_identifier)
#define isStatement(tok) ((tok) == tok_identifier || (tok) == tok_def || (tok) == tok_return || (tok) == tok_if || (tok) == tok_while || tok == tok_for)
#define isType(value, id) (value->getType()->getTypeID() == id)
//===----------------------------------------------------------------------===//
// Lexer
//===----------------------------------------------------------------------===//

// The lexer returns tokens [0-255] if it is an unknown character, otherwise one
// of these for known things.

#define IDENT_CHAR(A) (isalpha(A) || isdigit(A))

enum Token
{
  tok_invalid = 0, //0 is NULL char which won't be processed,so 0 is invalid
  tok_eof = -1,

  // commands
  tok_def = -2,
  tok_extern = -3,

  // primary
  tok_identifier = -4,
  tok_number = -5,
  tok_return = -6,
  tok_if = -7,
  tok_else = -8,
  tok_while = -9,
  tok_op = -10,
  tok_for = -11,
  tok_struct = -12,
};
enum Op_type
{
  op_add = 1,
  op_sub = 2,
  op_mul = 3,
  op_lt = 4,
  op_eq = 5,
  op_le = 6,
  op_ne = 7
};
static std::map<int, std::string> TypeValues; //map int to typeString
static std::map<std::string, int> OpValues;   //map opString to int
static FILE *fip;
static std::string IdentifierStr; // Filled in if tok_identifier || tok_number_string
static int NumValI;               // Filled in if tok_number_int
static double NumValD;            // Filled in if tok_number_double
static bool NumValB;              // Filled in if tok_number_bool
static char NumValC;              // Filled in if tok_number_char
static int ValType;               // Filled in if tok_number or tok_op

static void InitializeTypeValue()
{
  TypeValues[type_int] = "int";
  TypeValues[type_void] = "void";
  TypeValues[type_char] = "char";
  TypeValues[type_double] = "double";
  TypeValues[type_string] = "string";
  TypeValues[type_bool] = "bool";

  OpValues["+"] = op_add;
  OpValues["=="] = op_eq;
  OpValues["-"] = op_sub;
  OpValues["*"] = op_mul;
  OpValues["<"] = op_lt;
  OpValues["<="] = op_le;
  OpValues["!="] = op_ne;
}

/// gettok - Return the next token from standard input.
static int gettok()
{
  ///*****************************************

  ///use gettok implemented in lab1

  ///*****************************************
  static int LastChar = ' ';
  std::map<std::string, int>::iterator iter;

  // Skip any whitespace.
  while (isspace(LastChar))
    LastChar = fgetc(fip);

  if (isalpha(LastChar))
  {
    //------------------------------------------
    // TODO: identifier // definition // extern
    //-------------------------------------------
    IdentifierStr = LastChar;
    LastChar = fgetc(fip);
    while (IDENT_CHAR(LastChar))
    {
      IdentifierStr += LastChar;
      LastChar = fgetc(fip);
    }

    if (IdentifierStr == "int")
    {
      ValType = type_int;
      return tok_def;
    }
    if (IdentifierStr == "double")
    {
      ValType = type_double;
      return tok_def;
    }
    if (IdentifierStr == "bool")
    {
      ValType = type_bool;
      return tok_def;
    }
    if (IdentifierStr == "char")
    {
      ValType = type_char;
      return tok_def;
    }
    if (IdentifierStr == "string")
    {
      ValType = type_string;
      return tok_def;
    }
    if (IdentifierStr == "void")
    {
      ValType = type_void;
      return tok_def;
    }

    if (IdentifierStr == "true")
    {
      NumValB = true;
      ValType = type_bool;
      return tok_number;
    }
    if (IdentifierStr == "false")
    {
      NumValB = false;
      ValType = type_bool;
      return tok_number;
    }
    if (IdentifierStr == "extern")
      return tok_extern;
    if (IdentifierStr == "return")
      return tok_return;
    if (IdentifierStr == "if")
      return tok_if;
    if (IdentifierStr == "else")
      return tok_else;
    if (IdentifierStr == "while")
      return tok_while;
    if (IdentifierStr == "for")
      return tok_for;
    if (IdentifierStr == "struct")
      return tok_struct;
    return tok_identifier;
  }

  if (isdigit(LastChar))
  {
    //-----------------------------------
    // TODO: number_int // number_double
    //-----------------------------------
    int state = 0;
    int accept_tok = tok_number;
    ValType = type_int;
    IdentifierStr = std::string(1, (char)LastChar);
    LastChar = fgetc(fip);
    while (true)
    {
      switch (state)
      {
      case 0:
        if (isdigit(LastChar))
        {
          state = 0;
          IdentifierStr += LastChar;
          accept_tok = tok_number;
          ValType = type_int;
          LastChar = fgetc(fip);
        }
        else if (LastChar == '.')
        {
          state = 1;
          accept_tok = tok_invalid;
          LastChar = fgetc(fip);
        }
        else
        {
          state = 2;
        }
        break;
      case 1:
        if (isdigit(LastChar))
        {
          state = 1;
          if (accept_tok == tok_invalid)
          {
            IdentifierStr += ".";
          }
          IdentifierStr += LastChar;
          accept_tok = tok_number;
          ValType = type_double;
          LastChar = fgetc(fip);
        }
        else
        {
          state = 2;
        }
        break;
      case 2:
        if (accept_tok == tok_number && ValType == type_int)
        {
          NumValI = stoi(IdentifierStr);
        }
        else if (accept_tok == tok_number && ValType == type_double)
        {
          NumValD = stod(IdentifierStr);
        }
        else
        {
          LastChar = '.';
          NumValI = stoi(IdentifierStr);
          ValType = type_int;
          return tok_number;
        }
        return accept_tok;
        break;
      }
    }
  }

  if (LastChar == '\'')
  {
    LastChar = fgetc(fip);
    NumValC = LastChar;
    LastChar = fgetc(fip);
    if (LastChar != '\'')
      return tok_invalid;
    LastChar = fgetc(fip);
    ValType = type_char;
    return tok_number;
  }
  if (LastChar == '\"')
  {
    IdentifierStr = "";
    LastChar = fgetc(fip);
    while (LastChar != '\"' && LastChar != EOF)
    {
      IdentifierStr += LastChar;
      LastChar = fgetc(fip);
    }

    if (LastChar == EOF)
      return tok_invalid;
    LastChar = fgetc(fip);
    ValType = type_string;
    return tok_number;
  }

  if (LastChar == '+' || LastChar == '-' || LastChar == '*')
  {
    std::string op_type = "";
    op_type += (char)LastChar;
    ValType = OpValues[op_type];
    LastChar = fgetc(fip);
    return tok_op;
  }
  if (LastChar == '<')
  {
    std::string op_type = "";
    op_type += (char)LastChar;
    LastChar = fgetc(fip);
    if (LastChar != '=')
    {
      ValType = OpValues[op_type];
      return tok_op;
    }
    op_type += LastChar;
    ValType = OpValues[op_type];
    LastChar = fgetc(fip);
    return tok_op;
  }
  if (LastChar == '=' || LastChar == '!')
  {
    std::string op_type = "";
    op_type += (char)LastChar;
    int ThisChar = LastChar;
    LastChar = fgetc(fip);
    if (LastChar != '=')
      return ThisChar;
    op_type += LastChar;
    ValType = OpValues[op_type];
    LastChar = fgetc(fip);
    return tok_op;
  }
  // Check for end of file.
  if (LastChar == EOF)
    return tok_eof;

  // Otherwise, just return the character as its ascii value.
  int ThisChar = LastChar;
  LastChar = fgetc(fip);
  return ThisChar;
}

//===----------------------------------------------------------------------===//
// Abstract Syntax Tree (aka Parse Tree)
//===----------------------------------------------------------------------===//
// you don't have to modify this part. (of course it is ok to add something if needed)

//===----------------------------------------------------------------------===//
// Parser
//===----------------------------------------------------------------------===//

/// CurTok/getNextToken - Provide a simple token buffer.  CurTok is the current
/// token the parser is looking at.  getNextToken reads another token from the
/// lexer and updates CurTok with its results.
static int CurTok;
static int getNextToken() { return CurTok = gettok(); }

/// BinopPrecedence - This holds the precedence for each binary operator that is
/// defined.
static std::map<int, int> BinopPrecedence;
static std::vector<std::string> StrucTypes;
static bool foundStruct(std::string ident)
{
  std::vector<std::string>::iterator iter = std::find(StrucTypes.begin(), StrucTypes.end(), ident);
  return iter != StrucTypes.end();
}
/// LogError* - These are little helper functions for error handling.
/// you can add additional function to help you log error.
std::unique_ptr<ExprAST> LogError(const char *Str)
{
  fprintf(stderr, "Error: %s\n", Str);
  return nullptr;
}

std::unique_ptr<PrototypeAST> LogErrorP(const char *Str)
{
  LogError(Str);
  return nullptr;
}

std::unique_ptr<FunctionAST> LogErrorF(const char *Str)
{
  LogError(Str);
  return nullptr;
}
std::unique_ptr<StatementAST> LogErrorS(const char *Str)
{
  LogError(Str);
  return nullptr;
}
Value *LogErrorV(const char *Str)
{
  LogError(Str);
  return nullptr;
}
bool LogErrorB(const char *Str)
{
  LogError(Str);
  return false;
}
std::unique_ptr<StructAST> LogErrorStruct(const char *Str)
{
  LogError(Str);
  return nullptr;
}
///**********************************************

///use Parse*() function implemented in lab2

///**********************************************

static std::unique_ptr<ExprAST> ParseExpression();
static std::unique_ptr<StatementAST> ParseStatement();

/*TODO: Finish the Parse*() function to implement the Parser.
  We provide some implemented Parse* function for reference, like 
  ParseNumberExpr(), ParseExtern(), which are marked with "example", and you can use these functions directly. 
  >>>note: You can add some other Parse*() function to help you achieve this goal,
  >>>e.g. ParseParenExpr() which parenexpr ::= '(' expression ')'.
*/

/// numberexpr ::= number
/// example
static std::unique_ptr<ExprAST> ParseNumberExpr(int NumberType)
{
  std::unique_ptr<ExprAST> Result = nullptr;
  switch (NumberType)
  {
  case type_int:
    Result = std::make_unique<NumberIntExprAST>(NumValI);
    break;
  case type_double:
    Result = std::make_unique<NumberDoubleExprAST>(NumValD);
    break;
  case type_char:
    Result = std::make_unique<NumberCharExprAST>(NumValC);
    break;
  case type_bool:
    Result = std::make_unique<NumberBoolExprAST>(NumValB);
    break;
  case type_string:
    Result = std::make_unique<NumberStringExprAST>(IdentifierStr);
    break;
  default:
    return LogError("Unkonwn type literal");
  }
  getNextToken();
  return std::move(Result);
}

/// identifierexpr
/// <ident> or <callee>
/// TODO
static bool ParseArgs(std::vector<std::unique_ptr<ExprAST>> &Args)
{
  if (!isExpression(CurTok) && CurTok != ')')
    return LogErrorB("Expected <exp> or ')' in <callee>");
  while (true)
  {
    if (isExpression(CurTok))
    {
      auto exp = ParseExpression();
      if (exp)
        Args.push_back(std::move(exp));
    }
    else if (CurTok == ',')
    {
      getNextToken();
      if (!isExpression(CurTok))
        return LogErrorB("Expected <exr> in <callee>");
    }
    else
    {
      break;
    }
  }
  return true;
}
static std::unique_ptr<ExprAST> ParseIdentifierExpr()
{
  std::string ident = IdentifierStr;
  getNextToken(); //eat ident
  if (CurTok == '(')
  {
    std::vector<std::unique_ptr<ExprAST>> Args;
    getNextToken();
    if (!ParseArgs(Args))
      return nullptr;
    if (CurTok != ')')
      return LogError("Expected ')' in <callee>");
    getNextToken();
    return std::make_unique<CallExprAST>(ident, std::move(Args));
  }
  else if (CurTok == '.')
  {
    getNextToken();
    if (CurTok != tok_identifier)
      return LogError("Expected struct member name to load");
    std::string member = IdentifierStr;
    getNextToken();
    return std::make_unique<StructMemberReadAST>(ident, member);
  }
  else
  {
    return std::make_unique<VariableExprAST>(ident);
  }
}

/// expression
/// <exp>
/// TODO
static std::unique_ptr<ExprAST> HandleExpr(int precedence)
{
  std::unique_ptr<ExprAST> operate;
  switch (CurTok)
  {
  case '(':
    getNextToken();
    if (!isExpression(CurTok))
      return LogError("Expected expression in statement");
    operate = HandleExpr(0);
    if (CurTok != ')')
      return LogError("Expected ')' in statement");
    getNextToken();
    break;
  case tok_number:
    operate = ParseNumberExpr(ValType);
    break;
  case tok_identifier:
    operate = ParseIdentifierExpr();
    break;
  default:
    return LogError("Invalid expression in statement");
  }
  while (CurTok == tok_op)
  {
    int op_type = ValType;
    int rprecedence = BinopPrecedence[op_type];
    if (rprecedence <= precedence)
      break;
    getNextToken();
    auto right = HandleExpr(rprecedence);
    operate = std::make_unique<BinaryExprAST>(op_type, std::move(operate), std::move(right));
  }
  return operate;
}
static std::unique_ptr<ExprAST> ParseExpression()
{
  return HandleExpr(0);
}

/// statement
/// <stmt>
/// example
static std::unique_ptr<StatementAST> ParseAssign(std::string ident)
{
  if (CurTok == '(')
  {
    getNextToken();
    std::vector<std::unique_ptr<ExprAST>> Args;
    if (!ParseArgs(Args))
      return nullptr;
    if (CurTok != ')')
      return LogErrorS("Expected ')' in <callee>");
    getNextToken();
    auto expr = std::make_unique<CallExprAST>(ident, std::move(Args));
    return std::make_unique<NonAssignAST>(std::move(expr));
  }
  else if (CurTok == '.')
  {
    getNextToken();
    if (CurTok != tok_identifier)
      return LogErrorS("Expected member name to store");
    std::string member = IdentifierStr;
    getNextToken();
    if (CurTok != '=')
      return LogErrorS("Expected = to store member");
    getNextToken();
    auto expr = ParseExpression();
    return std::make_unique<StructMemberWriteAST>(ident, member, std::move(expr));
  }
  getNextToken();
  auto expr = ParseExpression();
  return std::make_unique<SimpAST>(ident, std::move(expr));
}
static std::unique_ptr<StatementAST> ParseSimp(std::string ident)
{

  auto assign = ParseAssign(ident);
  if (assign == nullptr)
    return nullptr;
  if (CurTok != ';')
    return LogErrorS("Expected ';' after statement");
  getNextToken();
  return assign;
}
static std::unique_ptr<StatementAST> ParseDecl(std::string type)
{
  if (type == TypeValues[type_void])
    return LogErrorS("Not surported void variable type");
  std::vector<std::string> variables;
  if (CurTok != tok_identifier)
    return LogErrorS("Expected identifier in declaration");
  variables.push_back(IdentifierStr);
  getNextToken();
  if (CurTok == '=')
  {
    getNextToken();
    auto expr = ParseExpression();
    if (CurTok != ';')
      return LogErrorS("Expected ';' after statement");
    getNextToken();
    return std::make_unique<DeclSimpAST>(type, variables[0], std::move(expr));
  }
  while (CurTok == ',')
  {
    getNextToken();
    if (CurTok != tok_identifier)
      return LogErrorS("Expected identifier in declaration");
    variables.push_back(IdentifierStr);
    getNextToken();
  }
  if (CurTok != ';')
    return LogErrorS("Expected ';' after statement");
  getNextToken();
  return std::make_unique<DeclAST>(type, std::move(variables));
}
static std::unique_ptr<StatementAST> ParseReturn()
{
  getNextToken();
  if (CurTok == ';')
  {
    getNextToken();
    return std::make_unique<ReturnAST>(nullptr);
  }
  auto expr = ParseExpression();
  if (CurTok != ';')
    return LogErrorS("Expected ';' after statement");
  getNextToken();
  return std::make_unique<ReturnAST>(std::move(expr));
}
static std::unique_ptr<BodyAST> ParseBlock()
{
  std::unique_ptr<BodyAST> body = std::make_unique<BodyAST>();
  if (CurTok == '{')
  {
    getNextToken();
    bool hasReturn = false;
    while (isStatement(CurTok))
    {
      int tok = CurTok;
      auto Stmt = ParseStatement();
      if (Stmt != nullptr && !hasReturn)
        body->add(std::move(Stmt));
      if (tok == tok_return)
        hasReturn = true;
    }
    if (CurTok != '}')
    {
      LogError("Expected '}' in block");
      return nullptr;
    }
    getNextToken();
  }
  else if (isStatement(CurTok))
  {
    auto Stmt = ParseStatement();
    if (Stmt == nullptr)
      return nullptr;
    body->add(std::move(Stmt));
  }
  else
  {
    LogError("Expected statement(s) in block");
    return nullptr;
  }
  return body;
}
static std::unique_ptr<StatementAST> ParseIF()
{
  getNextToken();
  if (CurTok != '(')
    return LogErrorS("Expected '(' for condition in if");
  getNextToken();
  auto cond = ParseExpression();
  if (cond == nullptr)
    return nullptr;
  if (CurTok != ')')
    return LogErrorS("Expected ')' for condition in if");
  getNextToken();
  auto trueBody = ParseBlock();
  if (trueBody == nullptr)
    return nullptr;
  if (CurTok == tok_else)
  {
    getNextToken();
    auto falseBody = ParseBlock();
    if (falseBody == nullptr)
      return nullptr;
    return std::make_unique<IFAST>(std::move(cond), std::move(trueBody), std::move(falseBody));
  }
  return std::make_unique<IFAST>(std::move(cond), std::move(trueBody), nullptr);
}
static std::unique_ptr<StatementAST> ParseWhile()
{
  getNextToken();
  if (CurTok != '(')
    return LogErrorS("Expected '(' for condition in if");
  getNextToken();
  auto cond = ParseExpression();
  if (cond == nullptr)
    return nullptr;
  if (CurTok != ')')
    return LogErrorS("Expected ')' for condition in if");
  getNextToken();
  auto body = ParseBlock();
  if (body == nullptr)
    return nullptr;
  return std::make_unique<WhileAST>(std::move(cond), std::move(body));
}
static std::unique_ptr<StatementAST> ParseFor()
{
  getNextToken();
  if (CurTok != '(')
    return LogErrorS("Expected '(' in for statement");
  getNextToken();
  std::unique_ptr<StatementAST> init;
  if (CurTok == tok_def)
  {
    int type = ValType;
    getNextToken();
    init = ParseDecl(TypeValues[type]);
    if (init == nullptr)
      return nullptr;
  }
  else if (CurTok == tok_identifier && foundStruct(IdentifierStr))
  {
    std::string ident = IdentifierStr;
    getNextToken();
    init = ParseDecl(ident);
    if (init == nullptr)
      return nullptr;
  }
  else if (CurTok == ';')
    getNextToken();
  else
    return LogErrorS("Expected a declare or ':' in for");
  std::unique_ptr<ExprAST> cond = ParseExpression();
  if (cond == nullptr)
    return nullptr;
  if (CurTok != ';')
    return LogErrorS("Expected ';' after contion in for");
  getNextToken();
  std::unique_ptr<StatementAST> after;
  if (CurTok == tok_identifier)
  {
    std::string ident = IdentifierStr;
    getNextToken();
    after = ParseAssign(ident);
    if (after == nullptr)
      return nullptr;
  }
  if (CurTok != ')')
    return LogErrorS("Expected ')' in for statement");
  getNextToken();
  auto block = ParseBlock();
  if (block == nullptr)
    return nullptr;
  return std::make_unique<ForAST>(std::move(init), std::move(cond), std::move(block), std::move(after));
}
static std::unique_ptr<StatementAST> ParseStatement()
{
  std::unique_ptr<StatementAST> ret;
  std::string ident = IdentifierStr;
  int type = ValType;
  switch (CurTok)
  {
  case tok_identifier:
    getNextToken();
    if (CurTok == '=' || CurTok == '(' || CurTok == '.')
      ret = ParseSimp(ident);
    else if (CurTok == tok_identifier)
    {
      if (!foundStruct(ident))
      {
        std::string err = "Unkown type : ";
        err += ident + " in declaration";
        return LogErrorS(err.c_str());
      }
      ret = ParseDecl(ident);
    }
    else
      return LogErrorS("Expected =,(,. or identifier in statement");
    break;
  case tok_def:
    getNextToken();
    ret = ParseDecl(TypeValues[type]);
    break;
  case tok_return:
    ret = ParseReturn();
    break;
  case tok_if:
    ret = ParseIF();
    break;
  case tok_while:
    ret = ParseWhile();
    break;
  case tok_for:
    ret = ParseFor();
    break;
  default:
    return LogErrorS("Expected statement in Block");
    break;
  }
  return ret;
}

/// prototype
/// <prototype>
/// an imcomplete parse function. It can parse <prototype> without args.
/// TODO
static std::unique_ptr<PrototypeAST> ParsePrototype()
{
  int IType = ValType;
  std::string FnType = (CurTok == tok_identifier) ? IdentifierStr : TypeValues[IType];
  getNextToken(); // eat ValType
  if (CurTok != tok_identifier)
    return LogErrorP("Expected function name in prototype");

  std::string FnName = IdentifierStr;
  getNextToken(); // eat FnName

  if (CurTok != '(')
    return LogErrorP("Expected '(' in prototype");
  getNextToken(); // eat '('.
  std::vector<std::string> ArgNames;
  std::vector<std::string> ArgTypes;
  // Maybe a FA is better here.
  if (CurTok != tok_def && CurTok != ')' && CurTok != tok_identifier)
    return LogErrorP("Expected arg type or ')' in prototype");
  while (true)
  {
    if (CurTok == tok_def || CurTok == tok_identifier)
    {
      if (CurTok == tok_identifier && !foundStruct(IdentifierStr))
      {
        std::string err = "Unknown type : ";
        err += IdentifierStr + " in define function args";
        return LogErrorP(err.c_str());
      }
      int argType = ValType;
      std::string type = (CurTok == tok_identifier) ? IdentifierStr : TypeValues[argType];
      getNextToken();
      if (CurTok != tok_identifier)
        return LogErrorP("Expected arg name in prototype");
      std::string argName = IdentifierStr;
      ArgNames.push_back(argName);
      ArgTypes.push_back(type);
      getNextToken();
      if (CurTok != ',' && CurTok != ')')
        return LogErrorP("Expected ',' or ')' in prototype");
    }
    else if (CurTok == ',')
    {
      getNextToken();
      if (CurTok != tok_def && CurTok != tok_identifier)
        return LogErrorP("Expected arg in prototype");
    }
    else
      break;
  }
  if (CurTok != ')')
    return LogErrorP("Expected ')' in prototype");

  // success.
  getNextToken(); // eat ')'.
  return std::make_unique<PrototypeAST>(FnName, std::move(ArgNames), std::move(ArgTypes), FnType);
}

/// definition ::= 'def' prototype expression
/// <function>
/// TODO
static std::unique_ptr<FunctionAST> ParseDefinition()
{
  auto Proto = ParsePrototype();
  if (CurTok != '{')
    return LogErrorF("Expected '{' in Function");
  getNextToken();
  auto body = std::make_unique<BodyAST>();
  bool hasReturn = false;
  while (isStatement(CurTok))
  {
    int tok = CurTok;
    auto Stmt = ParseStatement();
    if (Stmt != nullptr && !hasReturn)
      body->add(std::move(Stmt));
    if (tok == tok_return)
      hasReturn = true;
  }
  if (CurTok != '}')
    return LogErrorF("Expected '}' in Function");
  getNextToken();
  return std::make_unique<FunctionAST>(std::move(Proto), std::move(body));
}

/// external ::= 'extern' prototype
/// <gdecl>
/// example
static std::unique_ptr<PrototypeAST> ParseExtern()
{
  int isdef = getNextToken(); // eat extern.s
  if (isdef != tok_def && isdef != tok_identifier)
    return LogErrorP("Expected type declaration");
  auto Proto = ParsePrototype();
  if (CurTok != ';')
    return LogErrorP("Expected ';' in global declaration");
  getNextToken(); // eat ';'
  return Proto;
}
static std::unique_ptr<StructAST> ParseStruct()
{
  getNextToken();
  std::string name = IdentifierStr;
  std::vector<std::string> members, memberTypes;
  getNextToken();
  if (CurTok != '{')
    return LogErrorStruct("Expected '{' in struct defineation");
  getNextToken();
  while (CurTok != '}')
  {
    std::string typeStr;
    if (CurTok == tok_identifier)
    {
      if (!foundStruct(IdentifierStr))
      {
        std::string err = "Unkonwn type : ";
        err += IdentifierStr + " in declaration in struct";
        return LogErrorStruct(err.c_str());
      }
      typeStr = IdentifierStr;
    }
    else if (CurTok == tok_def)
    {
      typeStr = TypeValues[ValType];
    }
    else
      return LogErrorStruct("Expected declaration in struct");
    memberTypes.push_back(typeStr);
    getNextToken();
    if (CurTok != tok_identifier)
      return LogErrorStruct("Expected idetifier in declaration in struct");
    members.push_back(IdentifierStr);
    getNextToken();
    while (CurTok == ',')
    {
      getNextToken();
      if (CurTok != tok_identifier)
        return LogErrorStruct("Expected identifier in declaration");
      memberTypes.push_back(typeStr);
      members.push_back(IdentifierStr);
      getNextToken();
    }
    if (CurTok != ';')
      return LogErrorStruct("Expected ';' after declaration in Struct");
    getNextToken();
  }
  getNextToken();
  StrucTypes.push_back(name);
  return std::make_unique<StructAST>(name, std::move(members), std::move(memberTypes));
}

//===----------------------------------------------------------------------===//
// Code Generation
//===----------------------------------------------------------------------===//

static std::unique_ptr<LLVMContext> TheContext;
static std::unique_ptr<Module> TheModule;
static std::unique_ptr<IRBuilder<>> Builder;
static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
static std::unique_ptr<TypeSystem> TheTypeSystem;
static std::unique_ptr<ScopeContext> TheScopeContext;

static Function *CurrentFunction;
static std::string CurrentFnType;

Value *CastToBoolean(Value *cond)
{
  if (isType(cond, Type::IntegerTyID))
  {
    cond = Builder->CreateIntCast(cond, Type::getInt1Ty(*TheContext), true);
    return Builder->CreateICmpNE(cond, ConstantInt::get(Type::getInt1Ty(*TheContext), 0, true));
  }
  else if (isType(cond, Type::DoubleTyID))
    return Builder->CreateFCmpONE(cond, ConstantFP::get(*TheContext, APFloat(0.0)));
  else
    return nullptr;
}
// getFunction(Name) can return a Function structure variable, F, to caller, which can
// be used to creat a callee statement in codegen.
Function *getFunction(std::string Name)
{
  // First, see if the function has already been added to the current module.
  if (auto *F = TheModule->getFunction(Name))
    return F;

  // If not, check whether we can codegen the declaration from some existing
  // prototype.
  auto FI = FunctionProtos.find(Name);
  if (FI != FunctionProtos.end())
    return FI->second->codegen();

  // If no existing prototype exists, return null.
  return nullptr;
}

/*TODO: Finish the codegen() function to implement the Code Generation part.
  We provide some implemented codegen function for reference, like 
  NumberDoubleExprAST::codegen(), Value *VariableExprAST::codegen(), and you
  can use these functions directly. 
*/
Value *NumberIntExprAST::codegen()
{
  return ConstantInt::get(*TheContext, APInt(32, Val));
}
Value *NumberDoubleExprAST::codegen()
{
  return ConstantFP::get(*TheContext, APFloat(Val));
}
Value *NumberBoolExprAST::codegen()
{
  return Val ? ConstantInt::get(Type::getInt1Ty(*TheContext), 1, true) : ConstantInt::get(Type::getInt1Ty(*TheContext), 0, true);
}
Value *NumberCharExprAST::codegen()
{
  return ConstantInt::get(*TheContext, APInt(8, Val));
}
Value *NumberStringExprAST::codegen()
{
  return Builder->CreateGlobalString(Val, "string");
}

//example of codegen()
Value *VariableExprAST::codegen()
{
  // Look this variable up in the function.
  AllocaInst *alloca = TheScopeContext->getSymbolValue(Name);
  std::string errStr = "Unknown variable name: ";
  errStr += Name;
  if (!alloca)
    return LogErrorV(errStr.c_str());
  Value *V = Builder->CreateLoad(alloca->getAllocatedType(), alloca);
  return V;
}

//To Do
Value *BinaryExprAST::codegen()
{
  Value *L = LHS->codegen();
  Value *R = RHS->codegen();
  bool isDouble = L->getType()->isDoubleTy() || R->getType()->isDoubleTy();
  if (isDouble)
  {
    if (!L->getType()->isDoubleTy())
      L = Builder->CreateUIToFP(L, Type::getDoubleTy(*TheContext), "ftmp");
    if (!R->getType()->isDoubleTy())
      R = Builder->CreateUIToFP(R, Type::getDoubleTy(*TheContext), "ftmp");
    switch (Op)
    {
    case op_add:
      return Builder->CreateFAdd(L, R, "addtmp");
      break;
    case op_sub:
      return Builder->CreateFSub(L, R, "subtmp");
      break;
    case op_mul:
      return Builder->CreateFMul(L, R, "multmp");
      break;
    case op_lt:
      return Builder->CreateFCmpULT(L, R, "lttmp");
      break;
    case op_ne:
      return Builder->CreateFCmpONE(L, R, "netmp");
      break;
    case op_eq:
      return Builder->CreateFCmpOEQ(L, R, "eqtmp");
      break;
    default:
      return Builder->CreateFCmpOLE(L, R, "letmp");
    }
  }
  else
  {
    if (L->getType()->getTypeID() == TheTypeSystem->boolTy->getTypeID())
      L = Builder->CreateIntCast(L, Type::getInt32Ty(*TheContext), /*isSigned*/ false);
    if (R->getType()->getTypeID() == TheTypeSystem->boolTy->getTypeID())
      R = Builder->CreateIntCast(R, Type::getInt32Ty(*TheContext), false);
    switch (Op)
    {
    case op_add:
      return Builder->CreateAdd(L, R, "addtmp");
      break;
    case op_sub:
      return Builder->CreateSub(L, R, "subtmp");
      break;
    case op_mul:
      return Builder->CreateMul(L, R, "multmp");
      break;
    case op_lt:
      return Builder->CreateICmpSLT(L, R, "lttmp");
      break;
    case op_ne:
      return Builder->CreateICmpNE(L, R, "netmp");
      break;
    case op_eq:
      return Builder->CreateICmpEQ(L, R, "eqtmp");
      break;
    default:
      return Builder->CreateICmpSLE(L, R, "letmp");
    }
  }
  return nullptr;
}

//To Do
Value *CallExprAST::codegen()
{
  Function *func = getFunction(Callee);
  std::vector<Value *> putargs;
  for (auto &arg : Args)
    putargs.push_back(arg->codegen());
  ArrayRef<Value *> argsRef(putargs);
  return Builder->CreateCall(func, argsRef);
}

Value *StructMemberReadAST::codegen()
{
  auto varPtr = TheScopeContext->getSymbolValue(varName);
  auto structPtr = Builder->CreateLoad(varPtr, "structPtr");
  structPtr->setAlignment(4);
  if (!structPtr->getType()->isStructTy())
  {
    std::string err = "The variable ";
    err += varName + " is not struct";
    return LogErrorV(err.c_str());
  }
  std::string structName = structPtr->getType()->getStructName().str();
  long memberIndex = TheTypeSystem->getStructMemberIndex(structName, memberName);
  if (memberIndex < 0)
    return nullptr;
  std::vector<Value *> indices;
  indices.push_back(ConstantInt::get(TheTypeSystem->intTy, 0, false));
  indices.push_back(ConstantInt::get(TheTypeSystem->intTy, (uint64_t)memberIndex, false));
  auto ptr = Builder->CreateInBoundsGEP(varPtr, indices, "memberPtr");
  return Builder->CreateLoad(ptr);
}

// an imcomplete codegen function. It can generate IR for prototype whose types of args
// and return value are all double.
Function *PrototypeAST::codegen()
{
  // Make the function type:  double(double,double) etc.
  /*std::vector<Type *> Doubles(Args.size(), Type::getDoubleTy(*TheContext));
  FunctionType *FT =
      FunctionType::get(Type::getDoubleTy(*TheContext), Doubles, false);*/

  std::vector<Type *> types;
  for (auto argType : ArgTypes)
  {
    Type *type = TheTypeSystem->getVarType(argType);
    types.push_back(type);
  }
  ArrayRef<Type *> typeRef(types);
  Type *retType = TheTypeSystem->getVarType(FnType);
  FunctionType *FT = FunctionType::get(retType, typeRef, false);
  Function *F =
      Function::Create(FT, Function::ExternalLinkage, Name, TheModule.get());

  // Set names for all arguments.
  unsigned Idx = 0;
  for (auto &Arg : F->args())
    Arg.setName(Args[Idx++]);

  return F;
}

// an imcomplete codegen function.
// You should finish the ToDo part in this function
Function *FunctionAST::codegen()
{
  // Transfer ownership of the prototype to the FunctionProtos map, but keep a
  // reference to it for use below.
  auto &P = *Proto;
  FunctionProtos[Proto->getName()] = std::move(Proto);
  Function *TheFunction = getFunction(P.getName());
  if (!TheFunction)
    return nullptr;

  CurrentFunction = TheFunction;
  CurrentFnType = P.getReturnType();

  // Create a new basic block to start insertion into.
  BasicBlock *BB = BasicBlock::Create(*TheContext, "entry", TheFunction);
  Builder->SetInsertPoint(BB);
  TheScopeContext->pushBlock(BB);

  IRBuilder<> TmpB(BB, BB->begin());
  for (auto &Arg : TheFunction->args())
  {
    AllocaInst *alloca = TmpB.CreateAlloca(Arg.getType(), 0, Arg.getName());
    Builder->CreateStore(&Arg, alloca);
    TheScopeContext->setSymbol(std::string(Arg.getName()), alloca, P.getArgType(std::string(Arg.getName())), true);
  }

  if (Body->codegen())
  {
    //****************
    //ToDo: correctly create the RetVal use Builder.
    //****************
    // Validate the generated code, checking for consistency.
    bool isVoid = TheTypeSystem->getVarType(P.getReturnType())->getTypeID() == TheTypeSystem->voidTy->getTypeID();
    if (!isVoid && !TheScopeContext->hasReturn())
    {
      TheFunction->eraseFromParent();
      TheScopeContext->popBlock();
      std::string error = "Function: ";
      error += P.getName();
      error += ": Non void function must have return value";
      LogError(error.c_str());
      return nullptr;
    }
    if (isVoid && !TheScopeContext->hasReturn())
    {
      Builder->CreateRet(nullptr);
    }
    verifyFunction(*TheFunction);
    TheScopeContext->popBlock();
    return TheFunction;
  }

  // Error reading body, remove function.
  TheFunction->eraseFromParent();
  TheScopeContext->popBlock();
  return nullptr;
}
void createDecl(std::string type, std::string name)
{
  AllocaInst *Alloca;
  Value *tmp;
  BasicBlock *bb = TheScopeContext->currentBlock();
  IRBuilder<> TmpB(bb, bb->begin());
  Alloca = TmpB.CreateAlloca(TheTypeSystem->getVarType(type), 0, name.c_str());
  tmp = TheTypeSystem->getDefaultValue(type);
  if (tmp != nullptr)
    Builder->CreateStore(tmp, Alloca);
  TheScopeContext->setSymbol(name, Alloca, type, false);
}
bool DeclAST::codegen()
{
  for (auto name : variables)
  {
    createDecl(type, name);
  }
  return true;
}
bool SimpAST::codegen()
{
  AllocaInst *alloca = TheScopeContext->getSymbolValue(Name);
  if (!alloca)
  {
    LogError("Undefined variable");
    return false;
  }
  auto expr = expression->codegen();
  expr = TheTypeSystem->cast(expr, TheScopeContext->getSymbolType(Name), TheScopeContext->currentBlock());
  Builder->CreateStore(expr, alloca);
  return expr != nullptr;
}
bool DeclSimpAST::codegen()
{
  createDecl(type, Name);
  AllocaInst *alloca = TheScopeContext->getSymbolValue(Name);
  auto expr = expression->codegen();
  expr = TheTypeSystem->cast(expr, type, TheScopeContext->currentBlock());
  Builder->CreateStore(expr, alloca);
  return expr != nullptr;
}
bool ReturnAST::codegen()
{
  bool isVoid = CurrentFunction->getReturnType()->getTypeID() == Type::VoidTyID;
  if (isVoid)
  {
    if (expression != nullptr)
    {
      LogErrorV("return should not have value in void function.");
      TheScopeContext->setReturn(false);
      return false;
    }
    Builder->CreateRet(nullptr);
    TheScopeContext->setReturn(true);
    return true;
  }
  if (expression == nullptr)
  {
    LogErrorV("Non void function should have return value");
    TheScopeContext->setReturn(false);
    return false;
  }
  if (Value *RetVal = expression->codegen())
  {
    if(RetVal == nullptr)
      return false;
    RetVal = TheTypeSystem->cast(RetVal, CurrentFnType, TheScopeContext->currentBlock());
    if (RetVal == nullptr)
      return false;
    Builder->CreateRet(RetVal);
    // Validate the generated code, checking for consistency.
    TheScopeContext->setReturn(true);
    return true;
  }
  TheScopeContext->setReturn(false);
  return false;
}
bool NonAssignAST::codegen()
{
  auto expr = this->expression->codegen();
  return expr != nullptr;
}
bool BodyAST::codegen()
{
  for (auto &stmt : stmts)
    if (!stmt->codegen())
      return false;
  return true;
}
bool IFAST::codegen()
{
  Value *condValue = this->condExpr->codegen();
  if (!condValue)
    return false;
  condValue = CastToBoolean(condValue);
  if (!condValue)
    return false;
  bool trueHasReturn = false;
  bool falseHasReturn = false;
  BasicBlock *thenBB = BasicBlock::Create(*TheContext, "then", CurrentFunction);
  BasicBlock *falseBB = BasicBlock::Create(*TheContext, "else");
  BasicBlock *mergeBB = BasicBlock::Create(*TheContext, "ifcont");
  if (falseBlock != nullptr)
    Builder->CreateCondBr(condValue, thenBB, falseBB);
  else
    Builder->CreateCondBr(condValue, thenBB, mergeBB);
  Builder->SetInsertPoint(thenBB);
  TheScopeContext->pushBlock(thenBB);
  trueBlock->codegen();
  trueHasReturn = TheScopeContext->hasReturn();
  TheScopeContext->popBlock();
  if (!trueHasReturn)
  {
    Builder->CreateBr(mergeBB);
  }
  if (falseBlock != nullptr)
  {
    CurrentFunction->getBasicBlockList().push_back(falseBB);
    Builder->SetInsertPoint(falseBB);
    TheScopeContext->pushBlock(falseBB);
    falseBlock->codegen();
    falseHasReturn = TheScopeContext->hasReturn();
    TheScopeContext->popBlock();
    if (!falseHasReturn)
      Builder->CreateBr(mergeBB);
  }
  if (trueHasReturn && falseHasReturn)
    TheScopeContext->setReturn(true);
  CurrentFunction->getBasicBlockList().push_back(mergeBB);
  Builder->SetInsertPoint(mergeBB);
  TheScopeContext->setCurrentBlock(mergeBB);
  return true;
}
bool WhileAST::codegen()
{

  BasicBlock *condBB = BasicBlock::Create(*TheContext, "cond", CurrentFunction);
  BasicBlock *loopBB = BasicBlock::Create(*TheContext, "loop", CurrentFunction);
  BasicBlock *afterBB = BasicBlock::Create(*TheContext, "after", CurrentFunction);

  Builder->CreateBr(condBB);
  Builder->SetInsertPoint(condBB);
  Value *condValue = condExpr->codegen();
  if (!condValue)
    return false;
  condValue = CastToBoolean(condValue);
  if (!condValue)
    return false;
  Builder->CreateCondBr(condValue, loopBB, afterBB);
  Builder->SetInsertPoint(loopBB);
  TheScopeContext->pushBlock(loopBB);
  block->codegen();
  if (!TheScopeContext->hasReturn())
    Builder->CreateBr(condBB);
  TheScopeContext->popBlock();
  Builder->SetInsertPoint(afterBB);
  TheScopeContext->setCurrentBlock(afterBB);
  return true;
}
bool ForAST::codegen()
{
  BasicBlock *loopBB = BasicBlock::Create(*TheContext, "forloop", CurrentFunction);
  BasicBlock *afterBB = BasicBlock::Create(*TheContext, "forcont");
  TheScopeContext->pushBlock(TheScopeContext->currentBlock()); //Enter scope

  if (init != nullptr)
    init->codegen();
  Value *condValue = condExpr->codegen();
  if (!condValue)
    return false;
  condValue = CastToBoolean(condValue);
  if (!condValue)
    return false;
  Builder->CreateCondBr(condValue, loopBB, afterBB);
  Builder->SetInsertPoint(loopBB);
  TheScopeContext->setCurrentBlock(loopBB);
  block->codegen();
  if (after != nullptr)
    after->codegen();
  condValue = condExpr->codegen();
  condValue = CastToBoolean(condValue);
  if (!condValue)
    return false;
  Builder->CreateCondBr(condValue, loopBB, afterBB);

  TheScopeContext->popBlock(); //leave scope
  CurrentFunction->getBasicBlockList().push_back(afterBB);
  Builder->SetInsertPoint(afterBB);
  TheScopeContext->setCurrentBlock(afterBB);
  return true;
}
bool StructMemberWriteAST::codegen()
{
  auto varPtr = TheScopeContext->getSymbolValue(varName);
  auto structPtr = Builder->CreateLoad(varPtr, "structPtr");
  structPtr->setAlignment(4);
  if (!structPtr->getType()->isStructTy())
  {
    std::string err = "The variable ";
    err += varName + " is not struct";
    LogError(err.c_str());
    return false;
  }
  std::string structName = structPtr->getType()->getStructName().str();
  long memberIndex = TheTypeSystem->getStructMemberIndex(structName, memberName);
  if (memberIndex < 0)
    return false;
  std::vector<Value *> indices;
  indices.push_back(ConstantInt::get(TheTypeSystem->intTy, 0, false));
  indices.push_back(ConstantInt::get(TheTypeSystem->intTy, (uint64_t)memberIndex, false));
  auto value = expr->codegen();
  if (value == nullptr)
    return false;
  std::string memberType = TheTypeSystem->getStructMemberType(structName, memberName);
  value = TheTypeSystem->cast(value, memberType, TheScopeContext->currentBlock());
  if (value == nullptr)
    return false;
  auto ptr = Builder->CreateInBoundsGEP(varPtr, indices, "structMemberPtr");
  Builder->CreateStore(value, ptr);
  return true;
}
void StructAST::codegen()
{
  std::cout << "generate Struct " << name << ": " << std::endl;
  std::vector<Type *> typePointers;
  auto structType = StructType::create(*TheContext, name);
  TheTypeSystem->addStructType(name, structType);

  for (int i = 0; i < members.size(); i++)
  {
    std::cout << "    " << memberTypes[i] << "    " << members[i] << std::endl;
    TheTypeSystem->addStructMember(name, memberTypes[i], members[i]);
    typePointers.push_back(TheTypeSystem->getVarType(memberTypes[i]));
  }
  structType->setBody(typePointers);
}
//===----------------------------------------------------------------------===//
// Top-Level
//===----------------------------------------------------------------------===//
//don't modify this part

static void InitializeModuleAndPassManager()
{
  // Open a new context and module.
  TheContext = std::make_unique<LLVMContext>();
  TheModule = std::make_unique<Module>("my cool jit", *TheContext);

  // Create a new builder for the module.
  Builder = std::make_unique<IRBuilder<>>(*TheContext);
  TheTypeSystem = std::make_unique<TypeSystem>(*TheContext);
  TheScopeContext = std::make_unique<ScopeContext>();
}

static void HandleDefinition()
{
  if (auto FnAST = ParseDefinition())
  {
    if (auto *FnIR = FnAST->codegen())
    {
      fprintf(stderr, "Read function definition:");
      FnIR->print(errs());
      fprintf(stderr, "\n");
    }
  }
  else
  {
    // Skip token for error recovery.
    getNextToken();
  }
}

static void HandleExtern()
{
  if (auto ProtoAST = ParseExtern())
  {
    if (auto *FnIR = ProtoAST->codegen())
    {
      fprintf(stderr, "Read extern: ");
      FnIR->print(errs());
      fprintf(stderr, "\n");
      FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
    }
  }
  else
  {
    // Skip token for error recovery.
    getNextToken();
  }
}
static void HandleStruct()
{
  if (auto StruAST = ParseStruct())
  {
    fprintf(stderr, "Read struct: %s\n", StruAST->getName().c_str());
    StruAST->codegen();
  }
  else
    getNextToken();
}
/// top ::= definition | external | expression | ';'
static void MainLoop()
{
  while (true)
  {
    switch (CurTok)
    {
    case tok_eof:
      return;
    case tok_identifier:
    case tok_def:
      HandleDefinition();
      break;
    case tok_extern:
      HandleExtern();
      break;
    case tok_struct:
      HandleStruct();
      break;
    default:
      std::cout << "invalid input" << std::endl;
      getNextToken();
      break;
    }
  }
}
//===----------------------------------------------------------------------===//
// Main driver code.
//===----------------------------------------------------------------------===//
//don't modify this part

int main(int argc, char *argv[])
{
  if (argc < 2)
  {
    errs() << "You need to specify the file to compile";
    return 1;
  }
  char *FileName = argv[1];
  fip = fopen(FileName, "r");
  if (fip == nullptr)
  {
    errs() << "The file '" << FileName << "' is not existed";
    return 1;
  }

  InitializeTypeValue();

  BinopPrecedence[op_lt] = 10;
  BinopPrecedence[op_ne] = 10;
  BinopPrecedence[op_eq] = 10;
  BinopPrecedence[op_le] = 10;
  BinopPrecedence[op_add] = 20;
  BinopPrecedence[op_sub] = 20;
  BinopPrecedence[op_mul] = 40; // highest.
  getNextToken();

  InitializeModuleAndPassManager();
  MainLoop();

  InitializeAllTargetInfos();
  InitializeAllTargets();
  InitializeAllTargetMCs();
  InitializeAllAsmParsers();
  InitializeAllAsmPrinters();

  auto TargetTriple = sys::getDefaultTargetTriple();
  TheModule->setTargetTriple(TargetTriple);

  std::string Error;
  auto Target = TargetRegistry::lookupTarget(TargetTriple, Error);

  // Print an error and exit if we couldn't find the requested target.
  // This generally occurs if we've forgotten to initialise the
  // TargetRegistry or we have a bogus target triple.
  if (!Target)
  {
    errs() << Error;
    return 1;
  }

  auto CPU = "generic";
  auto Features = "";

  TargetOptions opt;
  auto RM = Optional<Reloc::Model>();
  auto TheTargetMachine =
      Target->createTargetMachine(TargetTriple, CPU, Features, opt, RM);
  TheModule->setDataLayout(TheTargetMachine->createDataLayout());

  auto Filename = "output.o";
  std::error_code EC;
  raw_fd_ostream dest(Filename, EC, sys::fs::OF_None);

  if (EC)
  {
    errs() << "Could not open file: " << EC.message();
    return 1;
  }

  legacy::PassManager pass;
  auto FileType = LLVMTargetMachine::CGFT_ObjectFile;

  if (TheTargetMachine->addPassesToEmitFile(pass, dest, nullptr, FileType))
  {
    errs() << "TheTargetMachine can't emit a file of this type";
    return 1;
  }

  pass.run(*TheModule);
  dest.flush();

  outs() << "Wrote " << Filename << "\n";

  return 0;
}
